from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any
from enum import Enum

class DatabaseType(str, Enum):
    MYSQL = "mysql"
    POSTGRESQL = "postgresql"
    SQLSERVER = "sqlserver"
    SQLITE = "sqlite"

class DatabaseConnection(BaseModel):
    db_type: DatabaseType = Field(..., description="数据库类型")
    host: Optional[str] = Field(None, description="主机地址")
    port: Optional[int] = Field(None, description="端口")
    database: str = Field(..., description="数据库名称")
    username: Optional[str] = Field(None, description="用户名")
    password: Optional[str] = Field(None, description="密码")
    charset: str = Field("utf8mb4", description="字符集")

class TableColumn(BaseModel):
    name: str
    data_type: str
    is_nullable: bool
    is_primary_key: bool
    is_foreign_key: bool
    default_value: Optional[str] = None
    comment: Optional[str] = None
    referenced_table: Optional[str] = None
    referenced_column: Optional[str] = None

class TableInfo(BaseModel):
    name: str
    comment: Optional[str] = None
    columns: List[TableColumn]
    primary_keys: List[str]
    foreign_keys: List[Dict[str, Any]]

class Relationship(BaseModel):
    from_table: str
    from_column: str
    to_table: str
    to_column: str
    relationship_type: str = "one-to-many"
    is_unique: bool = False  # 是否唯一约束，用于判断一对一关系
    is_many_to_many: bool = False  # 是否多对多关系
    cardinality: str = "1:N"  # 基数表示，如 1:1, 1:N, N:M
    description: Optional[str] = None  # 关系描述